import tensorflow as tf
import numpy as np
import pickle

tf.config.experimental_run_functions_eagerly(True)
tf.random.set_seed(0)
np.random.seed(0)

NEURONS = 50000
THRESHOLD = 0.8445

with open("data_set.pickle", 'rb') as f:
    data_set = pickle.load(f)
x_train = data_set["x_train"]
y_train = data_set["y_train"]
x_test = data_set["x_test"]
y_test = data_set["y_test"]


class ExtractDense(tf.keras.layers.Layer):
    def __init__(self, valid_len):
        super(ExtractDense, self).__init__()
        self.valid_len = valid_len

    def build(self, input_shape):
        self.w = self.add_weight("w",
                                 shape=[400, NEURONS],
                                 dtype=tf.double,
                                 initializer=tf.random_normal_initializer(),
                                 trainable=False)

    @staticmethod
    def normalization(x):
        max_value = tf.reduce_max(x, axis=0)
        min_value = tf.reduce_min(x, axis=0)
        x = (x - min_value) / (max_value - min_value)
        return x

    def connect(self, p):
        x = p[:, 0]
        y = p[:, 1]
        return np.concatenate([x, y])

    def compute_output_shape(self, input_shape):
        return input_shape[0], self.valid_len * 4

    def threshold(self, x, t):
        # tf.less返回两个张量各元素比较（x<y）得到的真假值组成的张量，如[False  True]，参见语句汇总
        cond = tf.less(x, tf.ones(tf.shape(x), dtype=x.dtype) * t)
        # tf.where有三个参数，根据第一个条件是否成立，当为True的时候选择第二个参数中的值，否则使用第三个参数中的值。参见语句汇总
        out = tf.where(cond, tf.zeros(tf.shape(x)), tf.ones(tf.shape(x)))
        out = tf.cast(out, dtype=np.int8)
        return out

    def call(self, inputs, **kwargs):
        image_num = inputs.get_shape()[0]
        code_all = []
        if image_num is not None:
            for image in inputs:
                p = tf.where(image)[:self.valid_len]
                x = p[:, 0]
                y = p[:, 1]
                x = self.normalization(x)
                y = self.normalization(y)
                p = tf.stack([x, y], axis=1)
                p = p.numpy()
                # 排序
                p_x = np.stack(sorted(p, key=lambda x_y: x_y[0]))
                p_y = np.stack(sorted(p, key=lambda x_y: x_y[1]))
                # 连接x和y
                p_x = self.connect(p_x)
                p_y = self.connect(p_y)
                # 连接codex,codey
                code = tf.convert_to_tensor(np.concatenate([p_x, p_y]))
                code_all.append(code)
            code_all = tf.stack(code_all)
            code_all = tf.matmul(code_all, self.w)
            # 阈值过滤
            code_all = self.threshold(code_all, THRESHOLD)
            return code_all
        else:
            inputs = tf.reshape(inputs, [-1, 56 * 56])
            inputs = inputs[:, :self.valid_len * 4]
            inputs = tf.cast(inputs, dtype=tf.double)
            inputs = tf.matmul(inputs, self.w)
            return inputs


model = tf.keras.models.Sequential([
    tf.keras.layers.InputLayer(input_shape=[56, 56]),
    ExtractDense(100),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.summary()

# 编译
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
# 训练模型
history = model.fit(x_train,
                    y_train,
                    epochs=5,
                    batch_size=64,
                    validation_data=(x_test, y_test))
np.save("translation.npy", history.history)
